import numpy as np
from keras.models import Sequential
from keras.layers import Dense, Activation
from keras.optimizers import SGD  # 随机梯度下降算法
from sklearn.datasets import load_iris  # IRIS数据集
from sklearn.preprocessing import LabelBinarizer  # 对数值进行标签化处理
from sklearn.model_selection import train_test_split

# 本例测试keras包的简单使用，使用python3.5环境
# keras默认使用TensorFlow作为后端，可以修改为Theano
# 修改配置路径为：C:\Users\Administrator\.keras\keras.json


def main():
    iris = load_iris()

    # 数据预处理
    train_data, test_data, train_traget, test_target = train_test_split(
        iris['data'],
        iris['target'],
        test_size=0.2,  # 测试数据占比
        random_state=1  # 随机分配
    )
    labels_train = LabelBinarizer().fit_transform(train_traget)
    labels_test = LabelBinarizer().fit_transform(test_target)

    # 构建模型
    model = Sequential(
        [
            Dense(5, input_dim=4),  # 5个输出，4个输入
            Activation('relu'),  # 激活函数为回路
            Dense(3),  # 3个输出
            Activation('sigmoid')  # 激活函数为sigmoid
        ]
    )

    # 优化器
    sgd = SGD(lr=0.01, decay=1e-6, momentum=0.9)
    model.compile(optimizer=sgd, loss="categorical_crossentropy")  # 指定优化器

    # 训练模型
    model.fit(train_data, labels_train, epochs=200, batch_size=40)  # 训练步进2000,每次训练40个

    # 测试模型
    print(model.predict_classes(test_data))

    # 保存训练结果
    model.sample_weights("./model")

    # 载入之前模型的训练结果
    model.load_weights("./model")


if __name__ == "__main__":
    main()
